/*  ----------------------------------------------------------------------------
 *  Arello Mobile
 *  Fast 3d Engine
 *  ----------------------------------------------------------------------------
 * Except where otherwise noted, this work is licensed under
 * http://creativecommons.org/licenses/by/3.0
 */

#include <vector>
#include "safe_new.h"

#ifdef SAFE_NEW

#undef new
#undef delete

#include <win32/win32.h> // For HFILE
#include <memory>
#include <SimpleList.h>
#include <ShlObj.h>

class RepFile
{
	HANDLE hfile;

	static bool get_special_folder(int shell_folder_id, char * buffer, int size)
	{
		LPITEMIDLIST list;
		if(S_OK != SHGetSpecialFolderLocation(0, shell_folder_id, &list))
			return false;

		//in docs regarding the SHGetPathFromIDList() function, there is no way to determine the size of the buffer
		if(!SHGetPathFromIDList(list, buffer))
			return false;

		return true;
	}

	static void get_app_name(char * buffer, int size)
	{
		char module_filename[4096];
		size_t i = GetModuleFileName(0, module_filename, 4096);
		size_t length = i;

		while(module_filename[i--] != '\\' && i >= 0)
		{}

		i += 2;	//skip slash and advance to next symbol
		while(i < length && i < size-1)
		{
			*buffer++ = module_filename[i++];
		}

		*buffer = 0;
	};

public:
	explicit RepFile(bool create)
	{
		char buffer[4096];
		get_special_folder(CSIDL_APPDATA, buffer, 4096);

		char app_name[4096];
		get_app_name(app_name, 4096);
		strcat(buffer, "/");
		strcat(buffer, app_name);
		strcat(buffer, "/memory.log");

		hfile = CreateFile( buffer, GENERIC_WRITE, FILE_SHARE_WRITE, NULL, create ? CREATE_ALWAYS : OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL | FILE_FLAG_RANDOM_ACCESS, NULL );
		if( hfile != INVALID_HANDLE_VALUE && !create )
		{
			SetFilePointer( hfile, 0, NULL, FILE_END );
		}
	}
	virtual ~RepFile()
	{
		if( hfile != INVALID_HANDLE_VALUE )
			CloseHandle( hfile );
	}
	void write( const char * str )const
	{
		if( hfile == INVALID_HANDLE_VALUE )
			return;
		DWORD read = 0;
		WriteFile( hfile, str, strlen(str), &read, NULL );
	}
};

struct MemNode : public SimpleListNode< MemNode >
{
	const char * file;
	unsigned line;
	unsigned size;
	unsigned magic_one;
};

const int MAGIC = ('g' << 24 ) + ('e' << 16) + ('r' << 8) + 'G';

char place_for_start[ sizeof( MemNode ) ];
char place_for_finish[ sizeof( MemNode ) ];

static MemNode * finish;
static MemNode * start;

static CRITICAL_SECTION critical_sec;

enum HeapState { NOT_CREATED, CREATED, DESTROYED };

static HeapState heap_state = NOT_CREATED;

static int simple_counter = 0;

static void destroy_heap()
{
	heap_state = DESTROYED;

	// Report here
	RepFile rep_file( false );
	rep_file.write( "Memory leaks:\n    addr |    size | line | file\n---------+---------+------+----------\n" );
	MemNode * curr = start->get_next();
	for( ; curr != finish; curr = curr->get_next() )
	{
		int * magic_two = (int*)(curr + sizeof(MemNode));

		char buff[1024];
		int bs = wsprintf( buff,"%8X |%8d |%5d | %s\n", reinterpret_cast<unsigned>( curr ) + sizeof(MemNode), curr->size, curr->line, curr->file );
		rep_file.write( buff );
	}

	DeleteCriticalSection( &critical_sec );
}

static void create_heap()
{
	if( heap_state == NOT_CREATED )
	{
		finish = new( place_for_finish ) MemNode;
		start = new( place_for_start ) MemNode;
		finish->link( start );

		RepFile rep_file( true );
		heap_state = CREATED;
		atexit( destroy_heap );
		InitializeCriticalSection( &critical_sec );
	}
}

void * __cdecl operator new( size_t s, const char * file, int line )
{
	create_heap();
	if( heap_state != CREATED )
		throw std::bad_alloc();

	EnterCriticalSection( &critical_sec );

	void * mama = malloc( s + 4 + sizeof(MemNode) );
	if( mama == 0 )
	{
		LeaveCriticalSection( &critical_sec );
		throw std::bad_alloc();
	}
	++simple_counter;

	char * chama = (char *)mama;
	int * magic_two = (int*)(chama + s + sizeof(MemNode));
	*magic_two = MAGIC;
	MemNode * just = new(mama) MemNode; // in-place new

	just->file = file;
	just->line = line;
	just->size = s;
	just->magic_one = MAGIC;

	just->link( start );

	void * p = (void *)( chama + sizeof(MemNode) );
	if( p == (void *)(0x7AE518) )
		p = p;

	LeaveCriticalSection( &critical_sec );

	return p;
}

void __cdecl operator delete( void * p, const char * file, int line )
{
	if( heap_state != CREATED )
		return;
	if( p == 0 )
		return;

	EnterCriticalSection( &critical_sec );
	
	char * chama = (char *)p - sizeof(MemNode);
	void * mama = chama;
	MemNode * just = (MemNode *)mama;
	int * magic_two = (int*)(chama + just->size + sizeof(MemNode));

	if( *magic_two != MAGIC || just->magic_one != MAGIC )
	{
		RepFile rep_file( false );
		char buff[1024];
		int bs = wsprintf( buff,"Magic Destroyed (%8X:%8X) %8X |%8d |%5d | %s\n",
					*magic_two, just->magic_one, 
					reinterpret_cast<unsigned>( just ) + sizeof(MemNode), 
					just->size, just->line, just->file );
		rep_file.write( buff );

		LeaveCriticalSection( &critical_sec );
		return;
	}
	just->magic_one = 0;
	*magic_two = 0;

	just->unlink();

	free( mama );
	--simple_counter;

	LeaveCriticalSection( &critical_sec );
}


void * __cdecl operator new( size_t s )
{
	return operator new( s, "<Runtime>", 0 );
}

void __cdecl operator delete( void * p )
{
	operator delete( p, "<Runtime>", 0 );
}

#endif

